import pandas as pd
from FlagEmbedding import BGEM3FlagModel
import os.path
from project_path import get_project_path
import numpy as np
from schema import db_connect


def batch_embedding(df: pd.DataFrame) -> pd.DataFrame:
    """批量生成嵌入"""
    # 匹配 question 列
    questions = df.filter(like='question').squeeze().tolist()
    # 生成嵌入结果
    embedding_model = BGEM3FlagModel(os.path.join(get_project_path(), 'llm', 'BAAI-bge-m3'))
    question_embedding = embedding_model.encode(questions, max_length=1024)['dense_vecs']
    df['question_embedding'] = list(question_embedding)
    # 保存结果
    df.to_pickle(os.path.join(get_project_path(), 'data', 'question_embedding.pkl'))
    return df


def single_embedding(question: str) -> np.ndarray:
    """逐一生成嵌入"""
    # 生成嵌入结果
    embedding_model = BGEM3FlagModel(os.path.join(get_project_path(), 'llm', 'BAAI-bge-m3'))
    return embedding_model.encode(question, max_length=1024)['dense_vecs']  # 'dense_vecs' 稠密向量


def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """计算余弦相似度"""
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))


if __name__ == '__main__':
    # 从数据库中加载测试问题
    test_ques_data = db_connect.execute_sql("SELECT test_question, test_num FROM test_question")
    test_ques_data = pd.DataFrame(list(test_ques_data.values())[0])
    batch_embedding(test_ques_data)
